import tkinter as tk
from tkinter import ttk, messagebox, filedialog
import requests
import json
import threading
import pyaudio
from pydub import AudioSegment
from pydub.playback import play
import tempfile
import os
import websockets
import asyncio
import numpy as np
import queue
import time
from OpusPlayer import play_opus_stream
import io
import wave

# 配置网络路径
HTTP_PATH = "http://192.168.2.4:8000"
WS_PATH = "ws://192.168.2.4:8000/realtime"
DEFAULT_VOICE = "ying"
DEFAULT_INPUT = "你不要让我做饭啦，我什么都能办到，但是真的不会做饭。"

# 音频格式设置
FORMAT_PCM = 0
FORMAT_OPUS = 1
FORMAT_WAV = 2

class VoiceSynthesizerApp:
    def __init__(self, root):
        self.root = root
        self.root.title("语音合成工具")
        self.root.geometry("800x700")
        self.current_voice = DEFAULT_VOICE
        self.audio = pyaudio.PyAudio()
        self.voice_map = {}
        self.voices_list = []
        self.stream_format = FORMAT_PCM
        
        # Websocket和音频队列状态
        self.realtime_active = False
        self.websocket = None
        self.realtime_thread = None
        self.stop_realtime_flag = False
        self.audio_queue = queue.Queue()
        self.playing = False
        self.play_thread = None
        self.stop_playback = threading.Event()
        
        # 主框架
        main_frame = ttk.Frame(root)
        main_frame.pack(fill='both', expand=True, padx=20, pady=20)
        
        # 音色管理区域
        voice_frame = ttk.LabelFrame(main_frame, text="音色管理")
        voice_frame.pack(fill='x', pady=(0, 15))
        
        # 音色选择和设置
        ttk.Label(voice_frame, text="当前音色:").grid(row=0, column=0, padx=5, pady=5, sticky='w')
        self.current_voice_label = ttk.Label(voice_frame, text=self.current_voice, width=15)
        self.current_voice_label.grid(row=0, column=1, padx=5, pady=5, sticky='w')
        
        ttk.Label(voice_frame, text="选择音色:").grid(row=0, column=2, padx=5, pady=5, sticky='w')
        self.voice_combobox = ttk.Combobox(voice_frame, width=15, state='readonly')
        self.voice_combobox.grid(row=0, column=3, padx=5, pady=5)
        self.voice_combobox.set(self.current_voice)
        self.voice_combobox.bind("<<ComboboxSelected>>", self.on_voice_selected)
        
        self.fetch_voices_btn = ttk.Button(voice_frame, text="获取可用音色", command=self.load_voices)
        self.fetch_voices_btn.grid(row=0, column=4, padx=5, pady=5)
        
        # 音色列表显示
        self.voices_text = tk.Text(voice_frame, height=6, width=80)
        self.voices_text.grid(row=1, column=0, columnspan=6, padx=10, pady=10)
        self.voices_text.insert('1.0', "点击'获取可用音色'按钮加载音色信息...")
        self.voices_text.config(state='disabled')
        
        # 公共输入区域
        input_frame = ttk.LabelFrame(main_frame, text="输入文本")
        input_frame.pack(fill='x', pady=(0, 15))
        
        self.common_input_text = tk.Text(input_frame, height=5)
        self.common_input_text.pack(fill='both', expand=True, padx=10, pady=5)
        self.common_input_text.insert('1.0', DEFAULT_INPUT)
        
        # 流式合成区域
        stream_frame = ttk.LabelFrame(main_frame, text="流式合成（实时播放支持）")
        stream_frame.pack(fill='x', pady=(0, 15))
        
        # 音频格式选择
        format_frame = ttk.Frame(stream_frame)
        format_frame.grid(row=0, column=0, columnspan=2, padx=5, pady=5, sticky='w')
        ttk.Label(format_frame, text="音频格式:").pack(side='left')
        self.format_var = tk.StringVar(value="PCM")
        self.format_combobox = ttk.Combobox(format_frame, textvariable=self.format_var, 
                                            values=["PCM", "Opus", "Wav"], width=10, state="readonly")
        self.format_combobox.pack(side='left', padx=5)
        self.format_combobox.bind("<<ComboboxSelected>>", self.on_format_selected)
        
        # 输出文件设置
        file_frame = ttk.Frame(stream_frame)
        file_frame.grid(row=1, column=0, columnspan=2, padx=5, pady=5, sticky='we')
        ttk.Label(file_frame, text="输出文件名:").pack(side='left')
        self.output_file = ttk.Entry(file_frame, width=25)
        self.output_file.insert(0, "output.pcm")
        self.output_file.pack(side='left', padx=5, fill='x', expand=True)
        
        # 操作按钮
        btn_frame = ttk.Frame(stream_frame)
        btn_frame.grid(row=2, column=0, columnspan=2, pady=10)
        
        self.play_stream_btn = ttk.Button(btn_frame, text="播放音频流", command=self.direct_stream_play)
        self.play_stream_btn.pack(side='left', padx=5)
        
        # 新增实时语音合成区域
        realtime_frame = ttk.LabelFrame(main_frame, text="实时语音合成 (WebSocket)")
        realtime_frame.pack(fill='x', pady=(15, 10))
        
        # 实时输入文本
        ttk.Label(realtime_frame, text="实时文本输入:").grid(row=0, column=0, padx=5, pady=5, sticky='w')
        self.realtime_input = tk.Text(realtime_frame, height=5)
        self.realtime_input.grid(row=0, column=1, padx=5, pady=5, sticky='we', columnspan=3)
        self.realtime_input.insert('1.0', "请在这里输入要实时合成的文本...")
        
        # 实时控制按钮
        realtime_btn_frame = ttk.Frame(realtime_frame)
        realtime_btn_frame.grid(row=1, column=0, columnspan=4, pady=10)
        
        self.realtime_start_btn = ttk.Button(realtime_btn_frame, text="连接WS", 
                                            command=self.start_realtime_synthesis)
        self.realtime_start_btn.pack(side='left', padx=5)
        
        self.realtime_send_btn = ttk.Button(realtime_btn_frame, text="发送文本", 
                                           command=self.send_realtime_text)
        self.realtime_send_btn.pack(side='left', padx=5)
        self.realtime_send_btn.config(state='disabled')
        
        self.realtime_stop_btn = ttk.Button(realtime_btn_frame, text="断开WS", 
                                           command=self.stop_realtime_synthesis)
        self.realtime_stop_btn.pack(side='left', padx=5)
        self.realtime_stop_btn.config(state='disabled')
        
        self.play_btn = ttk.Button(realtime_btn_frame, text="暂停/播放", command=self.toggle_playback)
        self.play_btn.pack(side='left', padx=5)
        self.play_btn.config(state='disabled')
        
        # 播放状态显示
        playback_status_frame = ttk.Frame(realtime_frame)
        playback_status_frame.grid(row=2, column=0, columnspan=4, pady=5)
        
        self.playback_status = ttk.Label(playback_status_frame, text="未开始播放")
        self.playback_status.pack()
        
        # 状态栏
        self.status_var = tk.StringVar()
        self.status_var.set("就绪")
        status_bar = ttk.Label(root, textvariable=self.status_var, relief=tk.SUNKEN, anchor=tk.W)
        status_bar.pack(side=tk.BOTTOM, fill=tk.X)
        
        # 默认加载音色
        self.load_voices()
    
    def on_voice_selected(self, event):
        selected_title = self.voice_combobox.get()
        if selected_title in self.voice_map:
            self.current_voice = self.voice_map[selected_title]
            self.current_voice_label.config(text=self.current_voice)
            self.set_status(f"当前音色: {selected_title}")
        else:
            self.set_status("未找到对应的音色")
    
    def on_format_selected(self, event):
        format_str = self.format_var.get()
        if format_str == "PCM":
            self.stream_format = FORMAT_PCM
            self.output_file.delete(0, 'end')
            self.output_file.insert(0, "output.pcm")
        else:  # Opus
            self.stream_format = FORMAT_OPUS
            self.output_file.delete(0, 'end')
            self.output_file.insert(0, "output.opus")
    
    def set_status(self, message):
        self.status_var.set(message)
        self.root.update_idletasks()
    
    def fetch_voices(self):
        try:
            self.set_status("正在获取音色列表...")
            self.fetch_voices_btn.config(state='disabled')
            response = requests.get(HTTP_PATH + "/voices")
            response.raise_for_status()
            data = response.json()
            
            self.voices_text.config(state='normal')
            self.voices_text.delete('1.0', 'end')
            self.voices_text.insert('1.0', json.dumps(data, indent=2, ensure_ascii=False))
            self.voices_text.config(state='disabled')
            
            # 更新音色下拉菜单和映射
            self.voice_map = {}
            voices_titles = []
            self.voices_list = data.get('voices', [])
            
            for voice in self.voices_list:
                title = voice.get('title', '')
                name = voice.get('name', '')
                if title and name:
                    self.voice_map[title] = name
                    voices_titles.append(title)
            
            if not voices_titles:
                voices_titles = [v.get('name', '') for v in self.voices_list]
                for voice in self.voices_list:
                    if 'name' in voice:
                        self.voice_map[voice['name']] = voice['name']
            
            self.voice_combobox['values'] = voices_titles
            if voices_titles:
                current_title = self.current_voice
                # 尝试找到当前音色对应的标题
                for voice in self.voices_list:
                    if voice.get('name') == self.current_voice:
                        current_title = voice.get('title', self.current_voice)
                        break
                if current_title in voices_titles:
                    self.voice_combobox.set(current_title)
                else:
                    self.voice_combobox.current(0)
                    self.current_voice = self.voice_map.get(voices_titles[0], '')
                    self.current_voice_label.config(text=self.current_voice)
            
            self.set_status(f"成功获取 {len(voices_titles)} 个音色")
        except requests.exceptions.RequestException as e:
            self.set_status(f"网络错误: {str(e)}")
            messagebox.showerror("错误", f"获取音色失败: {str(e)}")
        except Exception as e:
            self.set_status(f"错误: {str(e)}")
            messagebox.showerror("错误", f"获取音色失败: {str(e)}")
        finally:
            self.fetch_voices_btn.config(state='normal')
    
    def load_voices(self):
        threading.Thread(target=self.fetch_voices, daemon=True).start()
    
    def get_common_text(self):
        """获取公共输入框的文本内容"""
        return self.common_input_text.get("1.0", "end").strip()
    
    def get_realtime_text(self):
        """获取实时输入框的文本内容"""
        return self.realtime_input.get("1.0", "end-1c").strip()
    
    def direct_stream_play(self):
        text = self.get_common_text()
        if not text:
            messagebox.showwarning("警告", "请输入合成文本")
            return
        
        def run():
            try:
                self.set_status("正在实时播放音频流...")
                self.play_stream_btn.config(state='disabled')
                
                # 构建请求URL
                url = f"{HTTP_PATH}/stream?voice={self.current_voice}"
                if self.stream_format == FORMAT_OPUS:
                    url += "&format=opus"
                
                if self.stream_format == FORMAT_PCM:
                    self.play_pcm_stream(url, text)
                else:  # Opus格式
                    play_opus_stream(url, text)
                self.set_status("实时音频流播放完成")
            except requests.exceptions.RequestException as e:
                self.set_status(f"网络错误: {str(e)}")
                messagebox.showerror("错误", f"播放音频流失败: {str(e)}")
            except Exception as e:
                self.set_status(f"错误: {str(e)}")
                messagebox.showerror("错误", f"播放音频流失败: {str(e)}")
            finally:
                self.play_stream_btn.config(state='normal')
                self.set_status("就绪")
        
        threading.Thread(target=run, daemon=True).start()
    
    def play_wav(self):
        filename = self.filename.get().strip()
        
        def run():
            try:
                self.set_status("正在播放WAV音频...")
                self.play_wav_btn.config(state='disabled')
                
                # 使用pydub播放WAV文件
                audio = AudioSegment.from_file(filename, format="wav")
                play(audio)
                
                self.set_status("WAV音频播放完成")
            except Exception as e:
                self.set_status(f"播放错误: {str(e)}")
                messagebox.showerror("错误", f"播放WAV失败: {str(e)}")
            finally:
                self.play_wav_btn.config(state='normal')
        
        threading.Thread(target=run, daemon=True).start()
    
    def play_pcm_stream(self, url, text):
        # PCM格式使用PyAudio实时播放
        stream = self.audio.open(
            format=pyaudio.paFloat32,
            channels=1,
            rate=24000,
            output=True
        )
        
        with requests.post(
            url,
            headers={"Content-Type": "text/plain"},
            data=text.encode('utf-8'),
            stream=True
        ) as r:
            r.raise_for_status()
            
            # 获取并播放音频数据
            for data in r.iter_content(chunk_size=1024):
                if data:
                    stream.write(data)
        
        stream.stop_stream()
        stream.close()
                    
    # ===================== 实时语音合成功能 =====================
    
    def start_realtime_synthesis(self):
        """开始实时语音合成"""
        if self.realtime_active:
            self.set_status("实时合成已在进行中")
            return
            
        self.set_status("正在启动实时语音合成...")
        self.realtime_active = True
        self.stop_realtime_flag = False
        
        # 更新按钮状态
        self.realtime_start_btn.config(state='disabled')
        self.realtime_send_btn.config(state='normal')
        self.realtime_stop_btn.config(state='normal')
        self.play_btn.config(state='normal')
        
        # 清空音频队列
        while not self.audio_queue.empty():
            self.audio_queue.get()
        
        # 启动播放线程
        self.stop_playback.clear()
        if not self.play_thread or not self.play_thread.is_alive():
            self.play_thread = threading.Thread(target=self.audio_playback_worker, daemon=True)
            self.play_thread.start()
            self.playing = True
            self.playback_status.config(text="播放中")
        
        # 启动实时合成线程
        self.realtime_thread = threading.Thread(target=self.realtime_synthesis_worker, daemon=True)
        self.realtime_thread.start()
    
    def stop_realtime_synthesis(self):
        """停止实时语音合成"""
        self.set_status("正在停止实时合成...")
        self.stop_realtime_flag = True
        self.realtime_active = False
        
        # 更新按钮状态
        self.realtime_stop_btn.config(state='disabled')
        self.realtime_start_btn.config(state='normal')
        
        # 等待实时合成线程结束
        if self.realtime_thread and self.realtime_thread.is_alive():
            self.realtime_thread.join(timeout=2.0)
        
        # 停止播放
        self.stop_playback.set()
        self.playback_status.config(text="已停止")
        self.play_btn.config(state='disabled')
        self.realtime_send_btn.config(state='disabled')
    
    def send_realtime_text(self):
        """发送实时合成文本"""
        if not self.realtime_active:
            self.set_status("实时合成未启动")
            return
            
        text = self.get_realtime_text()
        if not text:
            self.set_status("请输入实时合成文本")
            return
            
        # 发送文本到WebSocket
        try:
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            loop.run_until_complete(self.send_text_to_websocket(text))
            self.set_status(f"已发送: {text[:30]}...")
        except Exception as e:
            self.set_status(f"发送失败: {str(e)}")
    
    def toggle_playback(self):
        """暂停/播放音频播放"""
        self.playing = not self.playing
        self.playback_status.config(text="播放中" if self.playing else "暂停")
    
    async def send_text_to_websocket(self, text):
        """异步发送文本到WebSocket"""
        if self.websocket:
            await self.websocket.send(text)
            self.set_status(f"已发送文本: {text[:50]}{'...' if len(text) > 50 else ''}")
        else:
            self.set_status("WebSocket连接未建立")
    
    def realtime_synthesis_worker(self):
        """实时合成工作线程（处理WebSocket连接）"""
        async def realtime_task():
            self.set_status("正在连接WebSocket服务器...")
            
            try:
                # 建立WebSocket连接
                self.websocket = await websockets.connect(WS_PATH, ping_interval=None, ping_timeout=None)
                self.set_status(f"已连接到 {WS_PATH}")
                
                # 发送配置信息
                config = {
                    "voice": self.current_voice,
                    "encoding": "pcm",  # 实时合成使用PCM格式
                    "sample_rate": 24000
                }
                await self.websocket.send(json.dumps(config))
                
                # 等待服务器的ready响应
                response = await self.websocket.recv()
                if isinstance(response, str):
                    resp_json = json.loads(response)
                    if resp_json.get('status') != 'ready':
                        self.set_status(f"服务器错误: {resp_json.get('message')}")
                        return
                    else:
                        self.set_status("服务器准备就绪，可以发送文本")
                
                # 持续接收音频数据
                self.set_status("开始接收音频数据...")
                
                while not self.stop_realtime_flag:
                    try:
                        # 接收音频数据
                        data = await asyncio.wait_for(self.websocket.recv(), timeout=0.5)
                        
                        if isinstance(data, bytes):
                            # 将接收到的数据放入队列供播放器使用
                            self.audio_queue.put(data)
                        elif isinstance(data, str):
                            try:
                                error_msg = json.loads(data)
                                self.set_status(f"服务器错误: {error_msg.get('message')}")
                            except:
                                self.set_status(f"收到服务器消息: {data}")
                    except asyncio.TimeoutError:
                        # 超时正常，用于检查停止标志
                        continue
                    except websockets.ConnectionClosed as e:
                        self.set_status(f"连接已关闭: {e}")
                        break
                
            except Exception as e:
                self.set_status(f"实时合成错误: {str(e)}")
            finally:
                # 关闭连接
                if self.websocket:
                    try:
                        await self.websocket.close()
                    except:
                        pass
                    self.websocket = None
                
                # 更新状态
                self.realtime_active = False
                self.stop_realtime_flag = False
                self.set_status("实时合成已停止")
                
                # 在UI线程中更新按钮状态
                self.root.after(0, lambda: [
                    self.realtime_start_btn.config(state='normal'),
                    self.realtime_stop_btn.config(state='disabled'),
                    self.realtime_send_btn.config(state='disabled'),
                    self.play_btn.config(state='disabled')
                ])
        
        # 运行异步任务
        asyncio.run(realtime_task())
    
    def audio_playback_worker(self):
        """音频播放工作线程（从队列播放音频数据）"""
        # 创建音频流
        stream = self.audio.open(
            format=pyaudio.paFloat32,
            channels=1,
            rate=24000,
            output=True,
            frames_per_buffer=1024
        )
        
        self.set_status("音频播放准备就绪")
        
        try:
            while not self.stop_playback.is_set():
                # 检查是否需要暂停
                if not self.playing:
                    time.sleep(0.1)
                    continue
                    
                # 从队列获取音频数据
                try:
                    data = self.audio_queue.get(timeout=0.5)
                    
                    # 播放音频
                    if data:
                        stream.write(data)
                    
                    self.audio_queue.task_done()
                    
                except queue.Empty:
                    # 队列为空是正常情况
                    continue
                except Exception as e:
                    self.set_status(f"播放错误: {str(e)}")
                    break
                    
        except Exception as e:
            self.set_status(f"播放线程错误: {str(e)}")
        finally:
            # 关闭音频流
            self.set_status("正在关闭音频流...")
            stream.stop_stream()
            stream.close()
            
            # 重置状态
            self.playing = False
            self.root.after(0, lambda: self.playback_status.config(text="播放线程已停止"))
            self.set_status("音频播放已停止")
    
    def __del__(self):
        """清理资源"""
        if self.audio:
            self.audio.terminate()

if __name__ == "__main__":
    root = tk.Tk()
    app = VoiceSynthesizerApp(root)
    root.mainloop()